

import numpy as np
def getRegions(genMap, regionSize):
    nLoci = len(genMap)
    # Genetic map is a matrix first column is chr, second column is something else.
    # We want to return a matrix with chromosome, start, stop.
    regions = []
    chrom = genMap[0]
    start = 0
    for i in range(nLoci):

        if chrom != genMap[i]:
            regions.append((chrom, start, i))
            start = i
            chrom = genMap[i]
        elif i - start >= regionSize:
            regions.append((chrom, start, i))
            start = i
    regions.append((chrom, start, nLoci))

    return np.array(regions, dtype = np.int64)

def getEffects(haplotypes, regions, qtls):
    effects = np.full((regions.shape[0], qtls.shape[0]), 0, dtype = np.float32)
    for i, region in enumerate(regions):
        chrom, start, stop = region
        effects[i,:] = np.sum(haplotypes[None, start:stop]*qtls[:, start:stop], 1)
    return effects

def imputeIndividualFromPhenotypes(ind, sire, dam, phenotypes, qtls, chrMap):
    print(ind.idx)
    nLoci = len(chrMap)
    regions = getRegions(chrMap, 100)


    patEffects = (getEffects(sire.haplotypes[0], regions, qtls), getEffects(sire.haplotypes[1], regions, qtls)) 
    matEffects = (getEffects(dam.haplotypes[0], regions, qtls), getEffects(dam.haplotypes[1], regions, qtls))

    patAssign, matAssign = fitAssignments(patEffects, matEffects, phenotypes, regions)
    ind.dosages = np.full(nLoci, 0, dtype = np.float32)
    for i, region in enumerate(regions):
        chrom, start, stop = region
        ind.dosages[start:stop] += (1-patAssign[i])*sire.haplotypes[0][start:stop]
        ind.dosages[start:stop] += patAssign[i]*sire.haplotypes[1][start:stop]

        ind.dosages[start:stop] += (1-matAssign[i])*dam.haplotypes[0][start:stop]
        ind.dosages[start:stop] += matAssign[i]*dam.haplotypes[1][start:stop]

def expNorm(vect):
    vect -= np.max(vect)
    tmp = np.exp(vect)
    return tmp/np.sum(tmp)

def fitAssignments(patEffects, matEffects, phenotypes, regions):
    nRegions = regions.shape[0]

    patOutput = np.full(nRegions, .5, dtype = np.float32)
    matOutput = np.full(nRegions, .5, dtype = np.float32)
    
    for i in range(1000):
        patOutput, matOutput = update(patOutput, matOutput, patEffects, matEffects, phenotypes, regions)

    return patOutput, matOutput


def update(patOutput, matOutput, patEffects, matEffects, phenotypes, regions):
    nRegions, nTraits = patEffects[0].shape

    phenotype_predicted = np.full(nTraits, 0, dtype = np.float32)
    
    phenotype_predicted += np.sum(patEffects[0] + patOutput[:,None]*(patEffects[1] - patEffects[0]), 0)
    phenotype_predicted += np.sum(matEffects[0] + matOutput[:,None]*(matEffects[1] - matEffects[0]), 0)

    scale = 0.001
    patUpdates = scale * 2 * np.sum((phenotypes - phenotype_predicted)[None,:]*(patEffects[1] - patEffects[0]), 1)
    matUpdates = scale * 2 * np.sum((phenotypes - phenotype_predicted)[None,:]*(matEffects[1] - matEffects[0]), 1)

    patOutput += patUpdates
    patOutput = np.maximum(np.minimum(patOutput, .999), .001)

    matOutput += matUpdates
    matOutput = np.maximum(np.minimum(matOutput, .999), .001)

    return patOutput, matOutput

